import os
import requests
import os
import json
import time
from datetime import datetime

class vtWrapper():
    def __init__(self,dataPath,savePath):
        self.originPath = dataPath
        self.finishedName = os.path.join(savePath,"finishedName.txt")
        self.VTJsonSaveFile = os.path.join(savePath,"lableJson.txt")
        self.vtLogPath = os.path.join(savePath,"vtlogs.txt")   
        if not os.path.exists(self.finishedName):
            with open(self.finishedName,"w") as file:
                file.close()
        if not os.path.exists(self.VTJsonSaveFile):
            with open(self.VTJsonSaveFile,"w") as file:
                file.close()
        if not os.path.exists(self.vtLogPath):
            with open(self.vtLogPath,"w") as file:
                file.close()

    def write_file(self,path,str1):
        with open(path,"a+") as f:
            f.write(str(datetime.now())+"\t"+str1+"\n")

    def get_sample_names(self):
        """
            获得需要处理的样本名以及路径
        """
        tar={}
        m,n=0,0
        for fileName in os.listdir(self.originPath):
            tar[fileName] = os.path.join(self.originPath,fileName)
        m=len(tar)
        with open(self.finishedName,"r") as file:
            for i in file.readlines():
                tempName=i.strip()
                if tempName in tar:
                    del tar[tempName]
        n=len(tar)
        print("[1]获得待处理样本路径：共{}个，还需处理{}个".format(m,n))
        return tar

    def get_json_from_VT(self,name,visitSpan):
        """
            将样本上传到vt获得json报告
        """ 
        url = "https://www.virustotal.com/api/v3/files"
        headers = {
            "accept": "application/json",
            "x-apikey": "b619a388252fb498b88e56e2efec3c0aa70be5edbb891494c5b14c4deb4e582c"
        }

        files = {"file": open(name, "rb")}
        self.write_file(self.vtLogPath,str(files))
        response = requests.post(url, files=files, headers=headers)
        self.write_file(self.vtLogPath,response.text)
        data=json.loads(response.text)
        analysisId=data["data"]["id"]
        self.write_file(self.vtLogPath, name+"---->analysis id = "+analysisId)
        # 查看分析情况，completed为分析完成

        # analyses不消耗quota
        url1 = "https://www.virustotal.com/api/v3/analyses/"+analysisId
        self.write_file(self.vtLogPath, "等待virusTotal处理样本")
        time.sleep(visitSpan)
        while(True):
            response = requests.get(url1, headers=headers)
            data=json.loads(response.text)
            sha256=data["meta"]["file_info"]["sha256"]
            is_avaliable=True if data["data"]["attributes"]["status"]=="completed" else False
            if(is_avaliable):
                self.write_file(self.vtLogPath,"样本分析完成")
                break
            self.write_file(self.vtLogPath,"---")
            time.sleep(visitSpan)

        #获取json报告
        url2 = "https://www.virustotal.com/api/v3/files/"+sha256
        response = requests.get(url2, headers=headers)
        data=json.loads(response.text)
        temp=json.dumps(data)

        with open(self.VTJsonSaveFile,"a+") as f:
            f.write(temp)
            f.write("\n")
        with open(self.finishedName,"a+") as f:
            f.write(name.split("/")[-1]+"\n")

    def analysisSamples(self):
        tar=self.get_sample_names()
        cnt=1
        for name in tar.keys():
            self.write_file(self.vtLogPath,"--------start analyse---------")
            self.write_file(self.vtLogPath,"第{}个样本进行分析".format(cnt))
            print("第{}个样本进行分析".format(cnt))
            # self.get_json_from_VT(tar[name],320)
            self.get_json_from_VT(tar[name],60)
            cnt+=1
            
if __name__=="__main__":
    originPath="/home/cyw/projects/malware_detected/originData/origin_customer_malware"
    savePath="/home/cyw/projects/malware_detected/saveData/origin_customer_malware_label"
    vtTools = vtWrapper(originPath,savePath)
    vtTools.analysisSamples()
